/*
* Copyright (C) 2020 ~ 2024 chaigec All Rights Reserved.
*
* Author:     chaigec <chaigec@163.com>
* Maintainer: chaigec <chaigec@163.com>

 * Use of this source code is governed by MIT license that can be found in the
 * LICENSE file in the root of the source tree. All contributing project authors
 * may be found in the AUTHORS file in the root of the source tree.
*/

#include "MkFlvMuxer.h"
#include "MkMedia/MkMediaUtil.h"
#include "MkAmf/MkAmf.h"

Uint32 CMkFlvMuxer::PacketHeader(const CMkVideoParam& VideoParam, const CMkAudioParam& AudioParam, const Uint32& FileSize, const Uint32& Duration, Uint8* pBuf, Uint32 nLen, CMkBufferList& ReBufferList)
{
    memset(pBuf, 0, nLen);
    Uint32 nOffset = 0;
    pBuf[0] = 'F';
    pBuf[1] = 'L';
    pBuf[2] = 'V';
    pBuf[3] = 0x01;
    pBuf[4] = ((MkCodecUnknown != VideoParam.GetCodec()) & 0x01);  // the eighth bit
    pBuf[4] |= ((MkCodecUnknown != AudioParam.GetCodec()) & 0x01) << 2;// the sixth bit
    pBuf[5] = 0x00;
    pBuf[6] = 0x00;
    pBuf[7] = 0x00;
    pBuf[8] = MkFlvHeaderSize & 0xFF;
    //packet pre tag size
    Uint8* pPreTag = pBuf + MkFlvHeaderSize;
    memset(pPreTag, 0, sizeof(Uint32));
    ReBufferList.Append(pBuf, MkFlvHeaderSize + sizeof(Uint32));
    nOffset += MkFlvHeaderSize + sizeof(Uint32);

    //packet script tag
    Uint32 nUsed = 0;
    PacketScriptTag(VideoParam, AudioParam, FileSize, Duration, pBuf + nOffset, nLen - nOffset, nUsed, ReBufferList);
    nOffset += nUsed;

    if (MkCodecUnknown != VideoParam.GetCodec()) {
        PacketVideoTagPreparation(VideoParam, pBuf + nOffset, nLen - nOffset, nUsed, ReBufferList);
        nOffset += nUsed;
    }
    if (MkCodecUnknown != AudioParam.GetCodec()) {
        PacketAudioTagPreparation(AudioParam, pBuf + nOffset, nLen - nOffset, nUsed, ReBufferList);
        nOffset += nUsed;
    }
    return NoneError;
}
//
//Uint32 CMkFlvMuxer::PacketEsPacket(const CMkVideoParam& VideoParam, const CMkAudioParam& AudioParam, BOOL bPacketSei, const MkEsPacket& EsPacket, Uint8* pBuf, Uint32 nLen, CMkBufferList& ReBufferList)
//{
//    Uint32 PreviousTagSize = 0;
//    Uint32 DataSize = EsPacket.nSize;
//    Uint8 *pBegin = pBuf + MkFlvTagHeaderSize;
//    Uint32 nOffset = MkFlvTagHeaderSize;
//    memset(pBuf, 0, nLen);
//    if (EsPacket.CodecType > MkCodecCavs) //audio
//    {
//        pBegin[0] = (CMkMediaUtil::CodecToRtmpAuidoType(EsPacket.CodecType) & 0x0F) << 4;
//        pBegin[0] |= 3 << 2;
//        pBegin[0] |= (AudioParam.GetSampleSize() & 0x01) << 1;
//        pBegin[0] |= AudioParam.GetStereo() & 0x01;
//        pBegin[1] = 0x01;
//        DataSize += 2;
//        PacketFlvTagHead(MkFlvTagAudio, DataSize, EsPacket.Dts, pBuf, ReBufferList);
//        nOffset += 2;
//        ReBufferList.Append(pBegin, 2);
//        ReBufferList.Append(EsPacket.pData, EsPacket.nSize);
//    } else {
//        int i = 0;
//        pBegin[i++] = ((MkH264NalSlice != EsPacket.NaluType ? 1 : 2) << 4) | (CMkMediaUtil::CodecToRtmpVideoType(EsPacket.CodecType) & 0x0F);
//        pBegin[i++] = 0 != EsPacket.nSize ? MkFlvAvcPacketNalu : MkFlvAvcPacketEndSequence;
//        Uint32 CompositionTime = EsPacket.Pts - EsPacket.Dts;
//        pBegin[i++] = (CompositionTime >> 16) & 0xFF;
//        pBegin[i++] = (CompositionTime >> 8) & 0xFF;
//        pBegin[i++] = CompositionTime & 0xFF;
//        //key frame
//        if (bPacketSei) {
//            MkString Sei = VideoParam.GetSei();
//            if (!Sei.empty()) {
//                bPacketSei = TRUE;
//                pBegin[i++] = (Sei.size() >> 24) & 0xFF;
//                pBegin[i++] = (Sei.size() >> 16) & 0xFF;
//                pBegin[i++] = (Sei.size() >> 8) & 0xFF;
//                pBegin[i++] = Sei.size() & 0xFF;
//                DataSize = i + Sei.size() + sizeof(Uint32) + EsPacket.nSize;
//                PacketFlvTagHead(MkFlvTagVideo, DataSize, EsPacket.Dts, pBuf, ReBufferList);
//                ReBufferList.Append(pBegin, i);
//                nOffset += i;
//                ReBufferList.Append(Sei.c_str(), Sei.size(), 0, TRUE);
//            } else {
//                bPacketSei = FALSE;
//            }
//        }
//        if (!bPacketSei) {
//            DataSize = i + EsPacket.nSize;
//            if (0 != EsPacket.nSize) {
//                DataSize += sizeof(Uint32);  //add size 4byte
//            }
//            PacketFlvTagHead(MkFlvTagVideo, DataSize, EsPacket.Dts, pBuf, ReBufferList);
//            ReBufferList.Append(pBegin, i);
//            nOffset += i;
//        }
//        if (0 != EsPacket.nSize) {
//            pBegin = pBuf + nOffset;
//            pBegin[0] = (EsPacket.nSize >> 24) & 0xFF;
//            pBegin[1] = (EsPacket.nSize >> 16) & 0xFF;
//            pBegin[2] = (EsPacket.nSize >> 8) & 0xFF;
//            pBegin[3] = EsPacket.nSize & 0xFF;
//            ReBufferList.Append(pBegin, sizeof(Uint32));
//            ReBufferList.Append(EsPacket.pData, EsPacket.nSize);
//            nOffset += sizeof(Uint32);
//        }
//    }
//    PreviousTagSize = DataSize + MkFlvTagHeaderSize;
//    pBegin = pBuf + nOffset;
//    pBegin[0] = (PreviousTagSize >> 24) & 0xFF;
//    pBegin[1] = (PreviousTagSize >> 16) & 0xFF;
//    pBegin[2] = (PreviousTagSize >> 8) & 0xFF;
//    pBegin[3] = PreviousTagSize & 0xFF;
//    ReBufferList.Append(pBegin, sizeof(Uint32));
//    return NoneError;
//}

Uint32 CMkFlvMuxer::PacketEsPacket(const CMkVideoParam& VideoParam, const CMkAudioParam& AudioParam, const MkEsPacket& EsPacket, Uint8* pBuf, Uint32 nLen, Uint32& nUsed, CMkBufferList& ReBufferList, BOOL bRtmp/* =FALSE*/)
{
    Uint8 NaluType = 0;
    Uint32 PreviousTagSize = 0;
    Uint8 *pBegin = pBuf + MkFlvTagHeaderSize;
    memset(pBuf, 0, nLen);
    int i = 0;
    CMkBufferList TmpList;
    if (EsPacket.CodecType > MkCodecCavs) //audio
    {
        pBegin[0] = (CMkMediaUtil::CodecToRtmpAuidoType(EsPacket.CodecType) & 0x0F) << 4;
        pBegin[0] |= 3 << 2;
        pBegin[0] |= (AudioParam.GetSampleSize() & 0x01) << 1;
        pBegin[0] |= AudioParam.GetStereo() & 0x01;
        pBegin[1] = 0x01;
        TmpList.Append(pBegin, 2);
        TmpList.Append(EsPacket.BufferList);
        pBegin += 2;
        nUsed += 2;
    } else {
        pBegin[i++] = ((EsPacket.bKeyPacket ? 1 : 2) << 4) | (CMkMediaUtil::CodecToRtmpVideoType(EsPacket.CodecType) & 0x0F);
        pBegin[i++] = EsPacket.BufferList.Empty() ? MkFlvAvcPacketEndSequence : MkFlvAvcPacketNalu;
        Uint32 CompositionTime = EsPacket.Pts - EsPacket.Dts;
        pBegin[i++] = (CompositionTime >> 16) & 0xFF;
        pBegin[i++] = (CompositionTime >> 8) & 0xFF;
        pBegin[i++] = CompositionTime & 0xFF;
        TmpList.Append(pBegin, i);
        pBegin += i;
        nUsed += i;
        i = 0;

        MkListBuffer::const_iterator it = EsPacket.BufferList.GetBufferList().begin();
        for (; it != EsPacket.BufferList.GetBufferList().end(); it++) {
            pBegin[i++] = (it->GetLength() >> 24) & 0xFF;
            pBegin[i++] = (it->GetLength() >> 16) & 0xFF;
            pBegin[i++] = (it->GetLength() >> 8) & 0xFF;
            pBegin[i++] = it->GetLength() & 0xFF;
            TmpList.Append(pBegin, i);
            TmpList.Append(*it);
            pBegin += i;
            nUsed += i;
            i = 0;
        }
    }
    if (!bRtmp) {
        PreviousTagSize = TmpList.GetTatolLen() + MkFlvTagHeaderSize;
        pBegin[i++] = (PreviousTagSize >> 24) & 0xFF;
        pBegin[i++] = (PreviousTagSize >> 16) & 0xFF;
        pBegin[i++] = (PreviousTagSize >> 8) & 0xFF;
        pBegin[i++] = PreviousTagSize & 0xFF;
        TmpList.Append(pBegin, i);
        nUsed += i + MkFlvTagHeaderSize;
        PacketFlvTagHead(EsPacket.CodecType > MkCodecCavs ? MkFlvTagAudio : MkFlvTagVideo, TmpList.GetTatolLen() - sizeof(Uint32), EsPacket.Dts, pBuf, ReBufferList);
        ReBufferList.Append(TmpList);
    } else {
        ReBufferList = TmpList;
        nUsed += MkFlvTagHeaderSize;
    }
    return NoneError;
}


Uint32 CMkFlvMuxer::PacketScriptTag(const CMkVideoParam& VideoParam, const CMkAudioParam& AudioParam, const Uint32& FileSize, const Uint32& Duration, Uint8* pBuf, Uint32 nLen, Uint32& nUsed, CMkBufferList& ReBufferList, BOOL bRtmp/* =FALSE*/, BOOL bRtmpServer/* = FALSE*/)
{
    Uint8 *pBegin = pBuf + MkFlvTagHeaderSize;  //offset flv tag header
    MkAmfObject Object;
    MkVector<MkAmfObject> VecObject;
    if (bRtmp && !bRtmpServer) {
        Object.strValue = "@setDataFrame";
        Object.DataType = MkAmfString;
        VecObject.push_back(Object);
    }

    Object.strValue = "onMetaData";
    Object.DataType = MkAmfString;
    VecObject.push_back(Object);

    Object.DataType = MkAmfEcmaArray;
    MkAmfObject ItemObject;
    if (0 != Duration) {
        ItemObject.DataType = MkAmfNumber;
        ItemObject.Name = "duration";
        ItemObject.NumberValue = Duration;
        Object.VecObjectValue.push_back(ItemObject);
    }

    if (MkCodecUnknown != VideoParam.GetCodec()) {
        ItemObject.DataType = MkAmfNumber;
        ItemObject.Name = "width";
        ItemObject.NumberValue = VideoParam.GetWidth();
        Object.VecObjectValue.push_back(ItemObject);

        ItemObject.DataType = MkAmfNumber;
        ItemObject.Name = "height";
        ItemObject.NumberValue = VideoParam.GetHeight();
        Object.VecObjectValue.push_back(ItemObject);

        if (0 != VideoParam.GetBitRate()) {
            ItemObject.DataType = MkAmfNumber;
            ItemObject.Name = "videodatarate";
            ItemObject.NumberValue = VideoParam.GetBitRate();
            Object.VecObjectValue.push_back(ItemObject);
        }

        ItemObject.DataType = MkAmfNumber;
        ItemObject.Name = "framerate";
        ItemObject.NumberValue = VideoParam.GetFrameRate();
        Object.VecObjectValue.push_back(ItemObject);

        ItemObject.DataType = MkAmfNumber;
        ItemObject.Name = "videocodecid";
        ItemObject.NumberValue = CMkMediaUtil::CodecToRtmpVideoType(VideoParam.GetCodec());
        Object.VecObjectValue.push_back(ItemObject);
    }

    if (MkCodecUnknown != AudioParam.GetCodec()) {
        if (0 != AudioParam.GetBitRate()) {
            ItemObject.DataType = MkAmfNumber;
            ItemObject.Name = "audiodatarate";
            ItemObject.NumberValue = AudioParam.GetBitRate();
            Object.VecObjectValue.push_back(ItemObject);
        }

        ItemObject.DataType = MkAmfNumber;
        ItemObject.Name = "audiosamplerate";
        ItemObject.NumberValue = AudioParam.GetSampleRate();
        Object.VecObjectValue.push_back(ItemObject);

        ItemObject.DataType = MkAmfNumber;
        ItemObject.Name = "audiosamplesize";
        ItemObject.NumberValue = AudioParam.GetSampleSize();
        Object.VecObjectValue.push_back(ItemObject);

        ItemObject.DataType = MkAmfBoolean;
        ItemObject.Name = "stereo";
        ItemObject.bValue = AudioParam.GetStereo();
        Object.VecObjectValue.push_back(ItemObject);

        ItemObject.DataType = MkAmfNumber;
        ItemObject.Name = "audiocodecid";
        ItemObject.NumberValue = CMkMediaUtil::CodecToRtmpAuidoType(AudioParam.GetCodec());
        Object.VecObjectValue.push_back(ItemObject);
    }

    if (0 != FileSize) {
        ItemObject.DataType = MkAmfNumber;
        ItemObject.Name = "filesize";
        ItemObject.NumberValue = FileSize;
        Object.VecObjectValue.push_back(ItemObject);
    }
    VecObject.push_back(Object);
    Uint8* pEnd = CMkAmf::EncodeObject(pBegin, pBuf + nLen, VecObject);  //offset flv tag header
    if (!pEnd) {
        return InvalidParams;
    }

    MkVector<MkAmfObject> VecParseObject;
    CMkAmf::DecodeObject(pBegin, pEnd, VecParseObject, FALSE);

    if (!bRtmp && !bRtmpServer) {
        Uint32 ContentLen = pEnd - pBegin;
        PacketFlvTagHead(MkFlvTagScript, ContentLen, 0, pBuf, ReBufferList);
        Uint32 PreviousTagSize = ContentLen + MkFlvTagHeaderSize;
        pEnd[0] = (PreviousTagSize >> 24) & 0xFF;
        pEnd[1] = (PreviousTagSize >> 16) & 0xFF;
        pEnd[2] = (PreviousTagSize >> 8) & 0xFF;
        pEnd[3] = PreviousTagSize & 0xFF;
        ReBufferList.Append(pBegin, ContentLen + sizeof(Uint32));
        nUsed = PreviousTagSize + sizeof(Uint32);
    } else {
        ReBufferList.Append(pBegin, pEnd - pBegin);
        nUsed = pEnd - pBegin + MkFlvTagHeaderSize;
    }
    return NoneError;
}

Uint32 CMkFlvMuxer::PacketFlvTagHead(const MkFlvTagType& TagType, const Uint32& DataSize, const Uint32& TimeStamp, Uint8* pBuf, CMkBufferList& ReBufferList)
{
    pBuf[0] = 0 << 4;  //always 0 Encrypted buffer
    pBuf[0] |= TagType & 0x1F;
    pBuf[1] = (DataSize >> 16) & 0xFF;
    pBuf[2] = (DataSize >> 8) & 0xFF;
    pBuf[3] = DataSize & 0xFF;
    pBuf[4] = (TimeStamp >> 16) & 0xFF;
    pBuf[5] = (TimeStamp >> 8) & 0xFF;
    pBuf[6] = TimeStamp & 0xFF;
    pBuf[7] = (TimeStamp >> 24) & 0xFF;
    pBuf[8] = 0x00;
    pBuf[9] = 0x00;
    pBuf[10] = 0x00;
    ReBufferList.Append(pBuf, MkFlvTagHeaderSize);
    return NoneError;
}

Uint32 CMkFlvMuxer::PacketVideoTagPreparation(const CMkVideoParam& VideoParam, Uint8* pBuf, Uint32 nLen, Uint32& nUsed, CMkBufferList& ReBufferList, BOOL bRtmp /*=FALSE*/)
{
    Uint8 *pBegin = pBuf + MkFlvTagHeaderSize; //offset flv tag header
    int i = 0;
    switch (VideoParam.GetCodec()) {
    case MkCodecH264:
    {
        Uint16 SpsLen = VideoParam.GetSps().size() & 0xFFFF;
        Uint16 PpsLen = VideoParam.GetPps().size() & 0xFFFF;
        pBegin[i++] = (1 << 4) | (MkVideoRtmpAvc & 0x0f);
        pBegin[i++] = MkFlvAvcPacketSequenceHeader;
        pBegin[i++] = 0x00;
        pBegin[i++] = 0x00;
        pBegin[i++] = 0x00;

        //pBegin[i++] = VideoParam.GetH264Param().ConfigurationVersion;
        //pBegin[i++] = VideoParam.GetH264Param().Profile;
        //pBegin[i++] = VideoParam.GetH264Param().Compatibility;
        //pBegin[i++] = VideoParam.GetH264Param().Level;
        pBegin[i++] = 0x01;
        pBegin[i++] = VideoParam.GetSps()[1];
        pBegin[i++] = VideoParam.GetSps()[2];
        pBegin[i++] = VideoParam.GetSps()[3];
        pBegin[i++] = 0xFF;
        //MkFree(pSps);

        pBegin[i++] = 0xe1;
        pBegin[i++] = SpsLen >> 8;
        pBegin[i++] = SpsLen & 0xFF;
        memcpy(&pBegin[i], VideoParam.GetSps().c_str(), SpsLen);
        i += SpsLen;
        pBegin[i++] = 0x01;
        pBegin[i++] = PpsLen >> 8;
        pBegin[i++] = PpsLen & 0xFF;
        memcpy(&pBegin[i], VideoParam.GetPps().c_str(), PpsLen);
        i += PpsLen;
    }
    break;
    case MkCodecH265:
    {
        Uint16 SpsLen = VideoParam.GetSps().size() & 0xFFFF;
        Uint16 PpsLen = VideoParam.GetPps().size() & 0xFFFF;
        Uint16 VpsLen = VideoParam.GetVps().size() & 0xFFFF;
        pBegin[i++] = (1 << 4) | (MkVideoRtmpHevC & 0x0f);
        pBegin[i++] = 0x00;
        i += 3;

        //#if 1
        pBegin[i++] = 0x01;
        pBegin[i++] = VideoParam.GetSps()[1];
        pBegin[i++] = VideoParam.GetSps()[2];
        pBegin[i++] = VideoParam.GetSps()[3];
        pBegin[i++] = VideoParam.GetSps()[4];
        pBegin[i++] = VideoParam.GetSps()[5];
        pBegin[i++] = VideoParam.GetSps()[6];
        pBegin[i++] = VideoParam.GetSps()[7];
        pBegin[i++] = VideoParam.GetSps()[8];
        pBegin[i++] = VideoParam.GetSps()[9];
        pBegin[i++] = VideoParam.GetSps()[10];
        pBegin[i++] = VideoParam.GetSps()[11];
        pBegin[i++] = VideoParam.GetSps()[12];
        i += 6;
        i += 2;
        pBegin[i++] = (0 << 6) | (0 << 3) | (0 << 2) | 3;
        pBegin[i++] = 0x03;

        //vps
        pBegin[i++] = 0xa0;
        pBegin[i++] = 0x00;
        pBegin[i++] = 0x01;
        pBegin[i++] = VpsLen >> 8;
        pBegin[i++] = VpsLen & 0xff;
        memcpy(&pBegin[i], VideoParam.GetVps().c_str(), VpsLen);
        i += VpsLen;

        /*sps*/
        pBegin[i++] = 0xa1;
        pBegin[i++] = 0x00;
        pBegin[i++] = 0x01;
        pBegin[i++] = SpsLen >> 8;
        pBegin[i++] = SpsLen & 0xFF;
        memcpy(&pBegin[i], VideoParam.GetSps().c_str(), SpsLen);
        i += SpsLen;

        /*pps*/
        pBegin[i++] = 0xa2;
        pBegin[i++] = 0x00;
        pBegin[i++] = 0x01;
        pBegin[i++] = PpsLen >> 8;
        pBegin[i++] = PpsLen & 0xFF;
        memcpy(&pBegin[i], VideoParam.GetPps().c_str(), PpsLen);
        i += PpsLen;
    }
    break;
    default:
        break;
    }
    if (!bRtmp) {
        PacketFlvTagHead(MkFlvTagVideo, i, 0, pBuf, ReBufferList);
        Uint32 PreviousTagSize = i + MkFlvTagHeaderSize;
        pBegin[i++] = (PreviousTagSize >> 24) & 0xFF;
        pBegin[i++] = (PreviousTagSize >> 16) & 0xFF;
        pBegin[i++] = (PreviousTagSize >> 8) & 0xFF;
        pBegin[i++] = PreviousTagSize & 0xFF;
        ReBufferList.Append(pBegin, i);
        nUsed = i + MkFlvTagHeaderSize;
    } else {
        ReBufferList.Append(pBegin, i);
        nUsed = i + MkFlvTagHeaderSize;
    }

    return NoneError;
}

Uint32 CMkFlvMuxer::PacketAudioTagPreparation(const CMkAudioParam& AudioParam, Uint8* pBuf, Uint32 nLen, Uint32& nUsed, CMkBufferList& ReBufferList, BOOL bRtmp /*=FALSE*/)
{
    Uint8 *pBegin = pBuf + MkFlvTagHeaderSize;  //offset flv tag
    Uint32 ContentLen = 0;
    if (MkCodecAac == AudioParam.GetCodec()) {
        pBegin[0] = (MkAudioRtmpAac & 0x0F) << 4;
        pBegin[0] |= 3 << 2;  //aac alway 3
        pBegin[0] |= (AudioParam.GetSampleSize() & 0x01) << 1;
        pBegin[0] |= AudioParam.GetStereo() & 0x01;
        pBegin[1] = 0x00;
        pBegin[2] = (AudioParam.GetObjectType() & 0x1F) << 3; //fisrt 5bit
        pBegin[2] |= (AudioParam.GetSampleRate() >> 1) & 0x07;
        pBegin[3] = AudioParam.GetSampleRate() & 0x01;
        pBegin[3] |= (AudioParam.GetChannel() & 0x0f) << 3;
        pBegin[3] |= (AudioParam.GetFrameLenghFlag() & 0x01) << 2;
        pBegin[3] |= (AudioParam.GetDependOnCoreCoder() & 0x01) << 1;
        pBegin[3] |= AudioParam.GetExternsionFlag() & 0x01;
        ContentLen = 4;
        if (!bRtmp) {
            PacketFlvTagHead(MkFlvTagAudio, ContentLen, 0, pBuf, ReBufferList);
            Uint32 PreviousTagSize = ContentLen + MkFlvTagHeaderSize;
            pBegin[4] = (PreviousTagSize >> 24) & 0xFF;
            pBegin[5] = (PreviousTagSize >> 16) & 0xFF;
            pBegin[6] = (PreviousTagSize >> 8) & 0xFF;
            pBegin[7] = PreviousTagSize & 0xFF;
            ReBufferList.Append(pBegin, 8);
            nUsed = 8;
        } else {
            ReBufferList.Append(pBegin, 4);
            nUsed = 4 + MkFlvTagHeaderSize;
        }
    }
    return NoneError;
}
